import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import javassist.*;
import org.jboss.windup.decompiler.fernflower.FernflowerDecompiler;
import org.jd.core.v1.ClassFileToJavaSourceDecompiler;
import org.jd.core.v1.api.loader.Loader;
import org.jd.core.v1.api.loader.LoaderException;
import org.jd.core.v1.api.printer.Printer;
import org.junit.Test;

import java.io.*;
import java.lang.instrument.Instrumentation;
import java.lang.reflect.Constructor;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.*;
import java.util.List;
import java.util.jar.JarFile;

public class MemShellCheck {
    private static IdentityHashMap<Thread, Thread> hooks;
    ObjectMapper objectMapper = new ObjectMapper();

    /**
     * 检测当前项目中是否存在shut hook函数，hook函数可能是木马植入的后台不死程序
     *
     * @throws ClassNotFoundException
     * @throws NoSuchFieldException
     * @throws NoSuchMethodException
     * @throws IllegalAccessException
     */
    @Test
    public void shutDownHookTest() throws ClassNotFoundException, NoSuchFieldException, NoSuchMethodException, IllegalAccessException {

        Runtime runtime = Runtime.getRuntime();
        Thread t1 = new Thread(() -> System.out.println("shutDown hook"));
        t1.setName("shutDown");
        runtime.addShutdownHook(t1);

        Class<?> clazz = Class.forName("java.lang.ApplicationShutdownHooks");
        Constructor<?> constructor = clazz.getDeclaredConstructor();
        constructor.setAccessible(true);

        Field field = clazz.getDeclaredField("hooks");
        field.setAccessible(true);
        Object obj = field.get(clazz);
        if (obj instanceof IdentityHashMap) {
            List<Thread> threadList = new ArrayList<>();
            ((IdentityHashMap<?, ?>) obj).forEach((k, v) -> {
                if (k instanceof Thread) {
                    Thread kThread = (Thread) k;
                    if (kThread.getName().toLowerCase(Locale.ROOT).contains("shutdown")) {
                        System.out.println("存在 shutdown hook,可能存在内存马");
                        System.out.println("尝试清除内存马");
                        threadList.add(kThread);
                        System.out.println("清除成功！");
                    }
                }
            });
            threadList.forEach(((IdentityHashMap<?, ?>) obj)::remove);

        }
    }

    /**
     * 检测当前路径下是否有多余的jar包
     */
    @Test
    public void agentFileInCurrentPahTest() {
        System.out.println(System.getProperty("user.dir"));
        String currentPath = System.getProperty("user.dir");
        Scanner in = new Scanner(System.in);
        File file = new File(currentPath);
        File[] filesName = file.listFiles();
        for (int i = 0; i < filesName.length; i++) {
            if (filesName[i].isFile()) {
                //路径
                System.out.println(filesName[i].getName());
                String fileName = filesName[i].getName();
                if (fileName.endsWith("jar") && !fileName.contains("css")) {
                    System.out.println("当前路径下存在未知的jar包,可能是内存马");
                }
            } else if (filesName[i].isDirectory()) {
                System.out.println("dir:" + filesName[i].getName());
            }
        }

    }

    /**
     * 检测存在不应该有的servlet
     */
    @Test
    public void servletTest() {
        String className = "org.apache.catalina.core.ApplicationFilterChain";
        //二进制对比,封禁url

    }

    /**
     * 检查是否有反弹shell存在
     */
    @Test
    public void shellCheck() {

    }

    /**
     * 检查是否有反弹shell存在
     */
    @Test
    public void urlRecord() {
        //代码中存在/bin/sh cmd.exe

    }

    /**
     *
     */
    public void codeCheck() {
        //检查是否存在 exec cmd  Runtime.getRuntime() 等代码字段
        System.out.println(System.getProperty("user.dir"));
        String currentPath = System.getProperty("user.dir");
        File file = new File(currentPath);
        File[] filesName = file.listFiles();

        Instrumentation instrumentation = null;
        Class<?>[] classes = instrumentation.getAllLoadedClasses();
        for (int i = 0; i < classes.length; i++) {
            System.out.println(classes[i].getName());
        }

    }

    @Test
    public void getClassFileTest() throws ClassNotFoundException, NoSuchAlgorithmException, IOException, NotFoundException, CannotCompileException {
        printMd5("org.apache.catalina.core.ApplicationFilterChain");
        ClassPool cp = ClassPool.getDefault();
        CtClass ctClass = cp.getCtClass("org.apache.catalina.core.ApplicationFilterChain");
        CtMethod m = ctClass.getDeclaredMethod("internalDoFilter");
        System.out.println(m.hashCode());

        m.insertBefore("System.out.println(\"dadsadsa\");");
        m.insertAfter("System.out.println(\"dadsadsfdsadsa\");");
        m.setName("dadsadsa");
        CtMethod m1 = ctClass.getDeclaredMethod("dadsadsa");
        System.out.println(m.getSignature());
        System.out.println(m1.hashCode());
        printMd5("org.apache.catalina.core.ApplicationFilterChain");

    }


    @Test
    public void printMd5Test() {
        try {
            printMd5("org.apache.catalina.core.ApplicationFilterChain", "internalDoFilter");
        } catch (ClassNotFoundException | NoSuchAlgorithmException | JsonProcessingException | NoSuchMethodException e) {
            e.printStackTrace();
        }
    }

    public void printMd5(String className) throws ClassNotFoundException, NoSuchAlgorithmException, JsonProcessingException {
        Class<?> clazz = Class.forName(className);
        byte[] clazzMd5s = MessageDigest.getInstance("md5").digest(objectMapper.writeValueAsBytes(clazz));
        System.out.println(Arrays.toString(clazzMd5s));
    }

    public void printMd5(String className, String method) throws ClassNotFoundException, NoSuchAlgorithmException, JsonProcessingException, NoSuchMethodException {
        Class<?> clazz = Class.forName(className);

        Method[] methods = clazz.getDeclaredMethods();
        for (Method m : methods) {
            System.out.println(m.getName());
            if (m.getName().equals("internalDoFilter")) {
                byte[] clazzMd5s = MessageDigest.getInstance("md5").digest(objectMapper.writeValueAsBytes(m));
                System.out.println(Arrays.toString(clazzMd5s));
            }
        }

    }

    @Test
    public void writeByteToFile() throws IOException {
        //将DemoServlet中的数据读取出来，写入另外一个DemoServlet文件中
        File file = new File("E:\\code\\servlet_demo\\target\\classes\\DemoServlet.class");
//        File file1 = new File("F:\\file\\DemoServlet.class");
        FileInputStream fileInputStream = new FileInputStream(file);

        BufferedReader reader = new BufferedReader(new FileReader(file));
        String str = null;
        StringBuilder sb = new StringBuilder();
        while ((str = reader.readLine()) != null) {
            sb.append(str);
        }

        System.out.println(sb);

        FileOutputStream fileOutputStream = new FileOutputStream("E:\\file\\DemoServlet.class");
        fileOutputStream.write(sb.toString().getBytes());

        reader.close();
        fileOutputStream.close();
        fileInputStream.close();


    }

    @Test
    public void deCompileTest() {
        FernflowerDecompiler fernflowerDecompiler = new FernflowerDecompiler();
//        Path path = new Window
//        fernflowerDecompiler.decompileClassFile(null,"F:\\code\\servlet_demo\\target\\classes\\DemoServlet.class","f:file/")
    }

    @Test
    public void loadTest() throws LoaderException, IOException {
//        File file = new File("e:\\code\\servlet_demo\\target\\classes\\DemoServlet.class");
//        FileInputStream fileInputStream = new FileInputStream(file);
//
//        BufferedReader reader = new BufferedReader(new FileReader(file));
//        String str = null;
//        StringBuilder sb = new StringBuilder();
//        while ((str = reader.readLine()) != null) {
//            sb.append(str);
//        }
//        FileOutputStream fileOutputStream = new FileOutputStream("E:/code/java_agent_asm/target/classes/DemoServlet.class");
//        fileOutputStream.write(sb.toString().getBytes());
//        reader.close();
//        fileOutputStream.close();
        InputStream is = this.getClass().getResourceAsStream("DemoServlet.class");
        if (is == null) {
            System.out.println("null");
        } else {
            try (InputStream in=is; ByteArrayOutputStream out=new ByteArrayOutputStream()) {
                byte[] buffer = new byte[1024];
                int read = in.read(buffer);

                while (read > 0) {
                    out.write(buffer, 0, read);
                    read = in.read(buffer);
                }

                System.out.println(Arrays.toString(out.toByteArray()));
            } catch (IOException e) {
                throw new LoaderException(e);
            }
        }
    }

    @Test
    public void jdTest() throws Exception {
        Loader loader = new Loader() {
            @Override
            public byte[] load(String internalName) throws LoaderException {
//
//                File file = new File("e:\\code\\servlet_demo\\target\\classes\\DemoServlet.class");
//                try {
//                    FileInputStream fileInputStream = new FileInputStream(file);
//                } catch (FileNotFoundException e) {
//                    e.printStackTrace();
//                }
//
//                BufferedReader reader = null;
//                try {
//                    reader = new BufferedReader(new FileReader(file));
//                } catch (FileNotFoundException e) {
//                    e.printStackTrace();
//                }
//                String str = null;
//                StringBuilder sb = new StringBuilder();
//                while (true) {
//                    try {
//                        assert reader != null;
//                        if ((str = reader.readLine()) == null) break;
//                    } catch (IOException e) {
//                        e.printStackTrace();
//                    }
//                    sb.append(str);
//                }
//                FileOutputStream fileOutputStream = null;
//                try {
//                    fileOutputStream = new FileOutputStream("E:/code/java_agent_asm/target/classes/DemoServlet.class");
//                } catch (FileNotFoundException e) {
//                    e.printStackTrace();
//                }
//                try {
//                    assert fileOutputStream != null;
//                    fileOutputStream.write(sb.toString().getBytes());
//                } catch (IOException e) {
//                    e.printStackTrace();
//                }
//                try {
//                    reader.close();
//                } catch (IOException e) {
//                    e.printStackTrace();
//                }
//                try {
//                    fileOutputStream.close();
//                } catch (IOException e) {
//                    e.printStackTrace();
//                }
                InputStream is = this.getClass().getResourceAsStream("DemoServlet.class");
                if (is == null) {
                    System.out.println("null");
                } else {
                    try (InputStream in=is; ByteArrayOutputStream out=new ByteArrayOutputStream()) {
                        byte[] buffer = new byte[1024];
                        int read = in.read(buffer);

                        while (read > 0) {
                            out.write(buffer, 0, read);
                            read = in.read(buffer);
                        }

                        System.out.println(Arrays.toString(out.toByteArray()));
                        return out.toByteArray();
                    } catch (IOException e) {
                        throw new LoaderException(e);
                    }
                }

//                InputStream is = this.getClass().getResourceAsStream(internalName + ".class");
//
//                if (is == null) {
//                    return null;
//                } else {
//                    try (InputStream in=is; ByteArrayOutputStream out=new ByteArrayOutputStream()) {
//                        byte[] buffer = new byte[1024];
//                        int read = in.read(buffer);
//
//                        while (read > 0) {
//                            out.write(buffer, 0, read);
//                            read = in.read(buffer);
//                        }
//
//                        return out.toByteArray();
//                    } catch (IOException e) {
//                        throw new LoaderException(e);
//                    }
//                }
                return null;
            }

            @Override
            public boolean canLoad(String internalName) {
//                return this.getClass().getResource("/" + internalName + ".class") != null;
                return true;
            }
        };
        Printer printer = new Printer() {
            private static final String TAB = "  ";
            private static final String NEWLINE = "\n";

            private int indentationCount = 0;
            private final StringBuilder sb = new StringBuilder();

            @Override
            public String toString() {
                return sb.toString();
            }

            @Override
            public void start(int maxLineNumber, int majorVersion, int minorVersion) {
            }

            @Override
            public void end() {
            }

            @Override
            public void printText(String text) {
                sb.append(text);
            }

            @Override
            public void printNumericConstant(String constant) {
                sb.append(constant);
            }

            @Override
            public void printStringConstant(String constant, String ownerInternalName) {
                sb.append(constant);
            }

            @Override
            public void printKeyword(String keyword) {
                sb.append(keyword);
            }

            @Override
            public void printDeclaration(int type, String internalTypeName, String name, String descriptor) {
                sb.append(name);
            }

            @Override
            public void printReference(int type, String internalTypeName, String name, String descriptor, String ownerInternalName) {
                sb.append(name);
            }

            @Override
            public void indent() {
                this.indentationCount++;
            }

            @Override
            public void unindent() {
                this.indentationCount--;
            }

            @Override
            public void startLine(int lineNumber) {
                for (int i = 0; i < indentationCount; i++) sb.append(TAB);
            }

            @Override
            public void endLine() {
                sb.append(NEWLINE);
            }

            @Override
            public void extraLine(int count) {
                while (count-- > 0) sb.append(NEWLINE);
            }

            @Override
            public void startMarker(int type) {
            }

            @Override
            public void endMarker(int type) {
            }
        };
        ClassFileToJavaSourceDecompiler decompiler = new ClassFileToJavaSourceDecompiler();
        decompiler.decompile(loader, printer, "DemoServlet");
        String source = printer.toString();
        System.out.println(source);

    }

    @Test
    public void getCurrentPath() {
        System.out.println(System.getProperty("user.dir"));
    }


    @Test
    public void writeFileToJar() throws IOException {
        File file = new File("servlet_demo-1.0-SNAPSHOT.jar");
        System.out.println();
        JarFile jarFile = new JarFile("servlet_demo-1.0-SNAPSHOT.jar");
        Enumeration<?> jarEntries = jarFile.entries();
//        jarEntries

    }

}
